package main

import (
	"flag"
	"fmt"
	"log"
	"net"
	"os"
	"os/signal"
	"sort"
	"sync"
	"syscall"
	"time"

	"github.com/google/gopacket"
	"github.com/google/gopacket/layers"
	"github.com/google/gopacket/pcap"
)

// UDPFlow 存储单个UDP流的信息
type UDPFlow struct {
	SrcIP   string
	SrcPort uint16
	DstIP   string
	DstPort uint16

	// 连接状态
	StartTime time.Time
	LastSeen  time.Time

	// 流量统计
	BytesSent       uint64
	BytesReceived   uint64
	PacketsSent     uint64
	PacketsReceived uint64
	AvgPacketSize   float64 // 平均包大小
	MinPacketSize   uint64  // 最小包大小
	MaxPacketSize   uint64  // 最大包大小

	// 应用识别相关
	AppName     string  // 应用名称
	AppCategory string  // 应用类别
	Confidence  float64 // 识别可信度(0-1)

	// Payload存储
	PayloadsSent     [][]byte // 源到目标的payload
	PayloadsReceived [][]byte // 目标到源的payload
	MaxPayloadsCount int      // 最大存储的payload数量
	MaxPayloadSize   int      // 单个payload的最大存储字节数
}

// FlowKey 用于标识唯一的UDP流
type FlowKey struct {
	SrcIP   string
	SrcPort uint16
	DstIP   string
	DstPort uint16
}

// 全局配置和状态
var (
	udpFlows          = make(map[FlowKey]*UDPFlow)
	flowMutex         = &sync.RWMutex{}
	iface             string
	snapshotLen       int32         = 1600
	promiscuous       bool          = false
	timeout           time.Duration = pcap.BlockForever
	printInterval     time.Duration = 5 * time.Second
	maxPayloadsCount  int           = 100            // 每个方向默认最多存100个payload
	maxPayloadSize    int           = 4096           // 单个payload默认最多存4096字节
	payloadStorageDir string        = "udp_payloads" // payload存储目录
)

func init() {
	flag.StringVar(&iface, "iface", "", "网络接口名称 (例如: eth0, en0)")
	flag.DurationVar(&printInterval, "interval", 5*time.Second, "统计信息打印间隔")
	flag.IntVar(&maxPayloadsCount, "max-payloads", 100, "每个方向最大存储的payload数量")
	flag.IntVar(&maxPayloadSize, "max-payload-size", 4096, "单个payload最大存储字节数")
	flag.StringVar(&payloadStorageDir, "payload-dir", "udp_payloads", "payload文件存储目录")
	flag.Parse()

	if iface == "" {
		log.Fatal("必须指定网络接口 (-iface)")
	}

	// 创建payload存储目录
	if err := os.MkdirAll(payloadStorageDir, 0755); err != nil {
		log.Fatalf("无法创建payload存储目录: %v", err)
	}
}

// 创建FlowKey，确保顺序一致
func getFlowKey(srcIP net.IP, srcPort uint16, dstIP net.IP, dstPort uint16) FlowKey {
	srcStr := fmt.Sprintf("%s:%d", srcIP, srcPort)
	dstStr := fmt.Sprintf("%s:%d", dstIP, dstPort)

	// 确保key的一致性，不管数据包方向
	if srcStr < dstStr {
		return FlowKey{
			SrcIP:   srcIP.String(),
			SrcPort: srcPort,
			DstIP:   dstIP.String(),
			DstPort: dstPort,
		}
	}
	return FlowKey{
		SrcIP:   dstIP.String(),
		SrcPort: dstPort,
		DstIP:   srcIP.String(),
		DstPort: srcPort,
	}
}

// 保存payload到内存（控制数量和大小）
func savePayload(payload []byte, payloads *[][]byte, maxCount, maxSize int) {
	// 限制单个payload大小
	if len(payload) > maxSize {
		payload = payload[:maxSize]
	}

	// 限制总数量（超过则移除最旧的）
	if len(*payloads) >= maxCount {
		*payloads = (*payloads)[1:] // 移除第一个元素（最旧的）
	}

	// 保存payload副本（避免引用原数据）
	payloadCopy := make([]byte, len(payload))
	copy(payloadCopy, payload)
	*payloads = append(*payloads, payloadCopy)
}

// 将内存中的payload保存到文件
func savePayloadsToFile(flow *UDPFlow) error {
	if len(flow.PayloadsSent) == 0 && len(flow.PayloadsReceived) == 0 {
		return nil // 没有payload需要保存
	}

	// 创建流专属目录
	flowDir := fmt.Sprintf("%s/%s_%d-%s_%d",
		payloadStorageDir,
		flow.SrcIP, flow.SrcPort,
		flow.DstIP, flow.DstPort)
	if err := os.MkdirAll(flowDir, 0755); err != nil {
		return err
	}

	// 保存发送的payload
	for i, payload := range flow.PayloadsSent {
		filename := fmt.Sprintf("%s/sent_%d_%d.bin",
			flowDir,
			flow.StartTime.UnixNano(),
			i)
		if err := os.WriteFile(filename, payload, 0644); err != nil {
			log.Printf("保存发送的payload失败: %v", err)
		}
	}

	// 保存接收的payload
	for i, payload := range flow.PayloadsReceived {
		filename := fmt.Sprintf("%s/received_%d_%d.bin",
			flowDir,
			flow.StartTime.UnixNano(),
			i)
		if err := os.WriteFile(filename, payload, 0644); err != nil {
			log.Printf("保存接收的payload失败: %v", err)
		}
	}

	return nil
}

// 解析DNS payload，提取域名
func parseDNS(payload []byte) (string, bool) {
	if len(payload) < 12 {
		return "", false
	}

	// DNS头部固定12字节，我们只简单检测查询域名
	// 更完整的实现应该使用gopacket的layers.DNS
	pos := 12
	for pos < len(payload) && pos < 512 { // 限制最大解析长度
		length := int(payload[pos])
		if length == 0 {
			break // 根域名
		}
		pos++
		if pos+length > len(payload) {
			break
		}
		// 简单判断是否是域名标签
		if length <= 63 { // 域名标签最大长度
			pos += length
		} else {
			break
		}
	}

	if pos+4 <= len(payload) {
		// 提取查询类型和类
		// 这里不做详细解析，只标记为DNS流量
		return "DNS查询", true
	}
	return "", false
}

// 基于端口识别应用
func identifyByPort(flow *UDPFlow) (appName string, appCategory string, confidence float64) {
	portAppMap := map[uint16][2]string{
		53:   {"DNS服务", "名称解析"},
		67:   {"DHCP服务器", "网络配置"},
		68:   {"DHCP客户端", "网络配置"},
		161:  {"SNMP服务", "网络管理"},
		123:  {"NTP服务", "时间同步"},
		5060: {"SIP协议", "语音通信"},
		5061: {"SIP TLS", "语音通信"},
		1900: {"SSDP协议", "设备发现"},
		4500: {"IPsec NAT-T", "虚拟专用网"},
		500:  {"ISAKMP", "虚拟专用网"},
		1194: {"OpenVPN", "虚拟专用网"},
		514:  {"Syslog", "系统日志"},
		3478: {"STUN", "网络穿透"},
		1935: {"RTMP", "视频流"},
		5004: {"RTP", "实时传输"},
		5005: {"RTCP", "实时控制"},
	}

	if appInfo, ok := portAppMap[flow.SrcPort]; ok {
		return appInfo[0], appInfo[1], 0.85
	}
	if appInfo, ok := portAppMap[flow.DstPort]; ok {
		return appInfo[0], appInfo[1], 0.85
	}

	// 动态端口范围判断
	if (flow.SrcPort >= 1024 && flow.SrcPort <= 65535) ||
		(flow.DstPort >= 1024 && flow.DstPort <= 65535) {
		return "动态UDP服务", "未知类别", 0.3
	}

	return "未知UDP服务", "未知类别", 0.2
}

// 基于payload内容识别应用
func identifyByPayload(flow *UDPFlow, payload []byte) (appName string, appCategory string, confidence float64) {
	if len(payload) == 0 {
		return "", "", 0
	}

	// 检测DNS
	if dnsQuery, ok := parseDNS(payload); ok {
		return dnsQuery, "名称解析", 0.9
	}

	// 检测NTP (简单特征)
	if len(payload) >= 4 && payload[0] == 0x1b {
		return "NTP协议", "时间同步", 0.8
	}

	// 检测SNMP (简单特征)
	if len(payload) >= 2 && payload[0] == 0x30 {
		return "SNMP协议", "网络管理", 0.7
	}

	// 检测RTP (简单特征)
	if len(payload) > 12 && (payload[1]&0x80) == 0x80 {
		return "RTP协议", "实时传输", 0.7
	}

	return "", "", 0
}

// 基于行为特征识别应用
func identifyByBehavior(flow *UDPFlow) (appCategory string, confidence float64) {
	duration := flow.LastSeen.Sub(flow.StartTime).Seconds()
	if duration < 1 {
		return "未知类别", 0.2
	}

	totalBytes := flow.BytesSent + flow.BytesReceived
	bytesPerSec := float64(totalBytes) / duration
	packetsPerSec := float64(flow.PacketsSent+flow.PacketsReceived) / duration

	switch {
	// 视频/语音流：高带宽、高频率
	case bytesPerSec > 1024*512 && packetsPerSec > 20:
		return "实时媒体流", 0.7
	// 游戏：中等带宽、高频率、小包
	case bytesPerSec < 1024*256 && packetsPerSec > 10 && flow.AvgPacketSize < 1024:
		return "在线游戏", 0.6
	// 即时通讯：低带宽、中等频率
	case bytesPerSec < 1024*64 && packetsPerSec > 2 && packetsPerSec < 50:
		return "即时通讯", 0.5
	// 批量数据传输：高带宽、低频率、大包
	case bytesPerSec > 1024*1024 && packetsPerSec < 10 && flow.AvgPacketSize > 1024*5:
		return "文件传输", 0.5
	}

	return "未知类别", 0.2
}

// 综合识别应用
func identifyApplication(flow *UDPFlow, payload []byte) {
	// 多维度识别
	portApp, portCat, portConf := identifyByPort(flow)
	payloadApp, payloadCat, payloadConf := identifyByPayload(flow, payload)
	behaviorCat, behaviorConf := identifyByBehavior(flow)

	// 融合结果（取可信度最高的）
	highestConf := 0.0
	bestAppName := "未知应用"
	bestCategory := "未知类别"

	if payloadConf > highestConf {
		highestConf = payloadConf
		bestAppName = payloadApp
		bestCategory = payloadCat
	}
	if portConf > highestConf {
		highestConf = portConf
		bestAppName = portApp
		bestCategory = portCat
	}
	if behaviorConf > highestConf {
		highestConf = behaviorConf
		bestCategory = behaviorCat
	}

	flow.AppName = bestAppName
	flow.AppCategory = bestCategory
	flow.Confidence = highestConf
}

// 处理UDP数据包
func handleUDPPacket(packet gopacket.Packet) {
	ipLayer := packet.Layer(layers.LayerTypeIPv4)
	if ipLayer == nil {
		return
	}
	ip, _ := ipLayer.(*layers.IPv4)

	udpLayer := packet.Layer(layers.LayerTypeUDP)
	if udpLayer == nil {
		return
	}
	udp, _ := udpLayer.(*layers.UDP)

	// 获取流密钥
	key := getFlowKey(ip.SrcIP, uint16(udp.SrcPort), ip.DstIP, uint16(udp.DstPort))
	isOriginalDirection := (key.SrcIP == ip.SrcIP.String() && key.SrcPort == uint16(udp.SrcPort))

	flowMutex.Lock()
	defer flowMutex.Unlock()

	// 检查流是否已存在，不存在则创建
	flow, exists := udpFlows[key]
	if !exists {
		flow = &UDPFlow{
			SrcIP:            key.SrcIP,
			SrcPort:          key.SrcPort,
			DstIP:            key.DstIP,
			DstPort:          key.DstPort,
			StartTime:        packet.Metadata().Timestamp,
			LastSeen:         packet.Metadata().Timestamp,
			MaxPayloadsCount: maxPayloadsCount,
			MaxPayloadSize:   maxPayloadSize,
			PayloadsSent:     make([][]byte, 0),
			PayloadsReceived: make([][]byte, 0),
		}
		udpFlows[key] = flow
	}

	// 更新最后活动时间
	flow.LastSeen = packet.Metadata().Timestamp

	// 更新流量统计
	packetSize := uint64(len(udp.Payload))
	if isOriginalDirection {
		flow.BytesSent += packetSize
		flow.PacketsSent++

		// 保存发送方向的payload
		if len(udp.Payload) > 0 {
			savePayload(udp.Payload, &flow.PayloadsSent, flow.MaxPayloadsCount, flow.MaxPayloadSize)
		}
	} else {
		flow.BytesReceived += packetSize
		flow.PacketsReceived++

		// 保存接收方向的payload
		if len(udp.Payload) > 0 {
			savePayload(udp.Payload, &flow.PayloadsReceived, flow.MaxPayloadsCount, flow.MaxPayloadSize)
		}
	}

	// 更新包大小统计
	if packetSize > 0 {
		if flow.MinPacketSize == 0 || packetSize < flow.MinPacketSize {
			flow.MinPacketSize = packetSize
		}
		if packetSize > flow.MaxPacketSize {
			flow.MaxPacketSize = packetSize
		}
	}

	// 计算平均包大小
	totalBytes := flow.BytesSent + flow.BytesReceived
	totalPackets := flow.PacketsSent + flow.PacketsReceived
	if totalPackets > 0 {
		flow.AvgPacketSize = float64(totalBytes) / float64(totalPackets)
	}

	// 识别应用
	identifyApplication(flow, udp.Payload)
}

// 打印统计信息
func printStats() {
	flowMutex.RLock()
	defer flowMutex.RUnlock()

	fmt.Printf("\n===== UDP监控统计 (%s) =====\n", time.Now().Format(time.RFC3339))
	fmt.Printf("总UDP流数量: %d\n", len(udpFlows))
	fmt.Println("------------------------------------------------------------------------------------------------------------------------------------------------------------------------")
	fmt.Printf("%-32s %-10s %-10s %-12s %-12s %-12s %-20s %-15s %-6.2s\n",
		"流", "上行(KB)", "下行(KB)", "平均包大小", "最小包大小", "最大包大小", "应用名称", "应用类别", "可信度")
	fmt.Println("------------------------------------------------------------------------------------------------------------------------------------------------------------------------")

	flows := make([]*UDPFlow, 0, len(udpFlows))
	for _, flow := range udpFlows {
		flows = append(flows, flow)
	}
	// 按最近活动时间排序
	sort.Slice(flows, func(i, j int) bool {
		return flows[i].LastSeen.After(flows[j].LastSeen)
	})

	// 只显示前20个流
	displayCount := 20
	if len(flows) < displayCount {
		displayCount = len(flows)
	}

	for i := 0; i < displayCount; i++ {
		flow := flows[i]
		flowStr := fmt.Sprintf("%s:%d -> %s:%d",
			flow.SrcIP, flow.SrcPort, flow.DstIP, flow.DstPort)
		upKB := float64(flow.BytesSent) / 1024
		downKB := float64(flow.BytesReceived) / 1024

		fmt.Printf("%-42s %-10.2f %-10.2f %-12.1f %-12d %-12d %-20s %-15s %-6.2f\n",
			flowStr,
			upKB,
			downKB,
			flow.AvgPacketSize,
			flow.MinPacketSize,
			flow.MaxPacketSize,
			flow.AppName,
			flow.AppCategory,
			flow.Confidence,
		)
	}

	fmt.Println("------------------------------------------------------------------------------------------------------------------------------------------------------------------------")
}

// 清理过期流（并将payload保存到文件）
func cleanupOldFlows() {
	flowMutex.Lock()
	defer flowMutex.Unlock()

	now := time.Now()
	// UDP流通常是短连接，设置较短的过期时间
	expirationTime := 2 * time.Minute

	for key, flow := range udpFlows {
		if now.Sub(flow.LastSeen) > expirationTime {
			// 清理前将payload保存到文件
			if err := savePayloadsToFile(flow); err != nil {
				log.Printf("保存flow payload失败: %v", err)
			}
			delete(udpFlows, key)
		}
	}
}

func main() {
	// 打开网络设备
	handle, err := pcap.OpenLive(iface, snapshotLen, promiscuous, timeout)
	if err != nil {
		log.Fatalf("无法打开设备 %s: %v", iface, err)
	}
	defer handle.Close()

	// 设置BPF过滤器，只捕获UDP数据包
	if err := handle.SetBPFFilter("udp"); err != nil {
		log.Fatalf("无法设置过滤器: %v", err)
	}

	// 启动统计打印goroutine
	go func() {
		ticker := time.NewTicker(printInterval)
		defer ticker.Stop()

		for range ticker.C {
			printStats()
			cleanupOldFlows()
		}
	}()

	// 捕获中断信号
	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)

	// 开始捕获数据包
	packetSource := gopacket.NewPacketSource(handle, handle.LinkType())
	packetChan := packetSource.Packets()

	fmt.Printf("开始监控UDP流 on %s...\n", iface)
	fmt.Printf("每 %v 打印一次统计信息\n", printInterval)
	fmt.Printf("Payload将存储在: %s\n", payloadStorageDir)
	fmt.Println("按Ctrl+C退出")

	for {
		select {
		case packet := <-packetChan:
			handleUDPPacket(packet)
		case <-sigChan:
			fmt.Println("\n收到退出信号，正在保存最后的统计信息和payload...")
			// 保存所有活跃流的payload
			flowMutex.RLock()
			for _, flow := range udpFlows {
				savePayloadsToFile(flow)
			}
			flowMutex.RUnlock()

			printStats()
			fmt.Println("监控已停止")
			return
		}
	}
}
